/*
 * Copyright 2008 Sebastian Köhler
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 * http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package calyriumdc.gui;

import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.Iterator;
import javax.swing.tree.DefaultMutableTreeNode;
import javax.swing.tree.DefaultTreeModel;
import javax.swing.tree.TreeNode;

public class SortedTreeModel extends DefaultTreeModel {
    
    private Comparator comp;

    public SortedTreeModel(TreeNode arg0, Comparator comp) {
        super(arg0);
        this.comp = comp;
    }
    
    public SortedTreeModel(TreeNode arg0, boolean arg1, Comparator comp) {
        super(arg0, arg1);
        this.comp = comp;
    }
    
    @Override
    public int getIndexOfChild(Object parent, Object child) {
        this.orderChildren(parent);
        return super.getIndexOfChild(parent, child);
    }

    @Override
    public Object getChild(Object parent, int index) {
        this.orderChildren(parent);
        return super.getChild(parent, index);
    }

    private void orderChildren(Object parent) {
        if (parent == null) {
            return;
        }
        if (parent instanceof DefaultMutableTreeNode) {
            DefaultMutableTreeNode node = (DefaultMutableTreeNode) parent;
            ArrayList children = Collections.list(node.children());
            Collections.sort(children, this.getComparator());
            node.removeAllChildren();
            Iterator childrenIterator = children.iterator();
            while (childrenIterator.hasNext()) {
                node.add((DefaultMutableTreeNode) childrenIterator.next());
            }
        }
    }

    public Comparator getComparator() {
        return comp;
    }
    
}
